////////////////////////////////////////////////////////////////////////////////////////////////////
//
//  Project:  Embedded Learning Library (ELL)
//  File:     DSPNodesTiming.cpp (nodes_test)
//  Authors:  Chuck Jacobs, Byron Changuion, Kern Handa
//
////////////////////////////////////////////////////////////////////////////////////////////////////

#include "DSPNodesTiming.h"
#include "NodesTestUtilities.h"

#include <dsp/include/Convolution.h>
#include <dsp/include/WinogradConvolution.h>

#include <math/include/MathConstants.h>
#include <math/include/Tensor.h>
#include <math/include/TensorOperations.h>

#include <model/include/InputNode.h>
#include <model/include/Model.h>
#include <model/include/Node.h>

#include <nodes/include/DiagonalConvolutionNode.h>
#include <nodes/include/SimpleConvolutionNode.h>
#include <nodes/include/UnrolledConvolutionNode.h>
#include <nodes/include/WinogradConvolutionNode.h>

#include <predictors/neural/include/ConvolutionalLayer.h>

#include <testing/include/testing.h>

#include <utilities/include/Exception.h>
#include <utilities/include/MillisecondTimer.h>
#include <utilities/include/RandomEngines.h>

#include <cmath>
#include <iostream>
#include <numeric>
#include <sstream>
#include <string>

using namespace ell;
using namespace nodes;
using namespace std::string_literals;

//
// Helpers
//
namespace
{
template <typename ValueType>
std::ostream& operator<<(std::ostream& os, const std::vector<ValueType>& vec)
{
    os << "[";
    for (auto x : vec)
    {
        os << x << " ";
    }
    os << "]";
    return os;
}

template <typename ElementType>
void FillRandomVector(std::vector<ElementType>& vector, ElementType min = -1, ElementType max = 1)
{
    auto randomEngine = utilities::GetRandomEngine("123");
    std::uniform_real_distribution<ElementType> uniform(min, max);
    auto rand = [&randomEngine, &uniform]() { return uniform(randomEngine); };
    std::generate(vector.begin(), vector.end(), rand);
}

template <typename ElementType>
void FillDataVector(std::vector<ElementType>& vector, int numRows, int numColumns, int numChannels)
{
    int vectorIndex = 0;
    for (int rowIndex = 0; rowIndex < numRows; ++rowIndex)
    {
        for (int columnIndex = 0; columnIndex < numColumns; ++columnIndex)
        {
            for (int channelIndex = 0; channelIndex < numChannels; ++channelIndex)
            {
                ElementType rowValue = rowIndex < 3 && columnIndex < 2 ? 2 * rowIndex : 0;
                ElementType columnValue = rowIndex < 3 && columnIndex < 2 ? columnIndex + 1 : 0;
                vector[vectorIndex++] = rowValue + columnValue;
            }
        }
    }
}

std::string GetConvAlgName(dsp::ConvolutionMethodOption alg)
{
    switch (alg)
    {
    case dsp::ConvolutionMethodOption::automatic:
        return "automatic";
    case dsp::ConvolutionMethodOption::simple:
        return "simple";
    case dsp::ConvolutionMethodOption::unrolled:
        return "unrolled";
    case dsp::ConvolutionMethodOption::diagonal:
        return "diagonal";
    case dsp::ConvolutionMethodOption::winograd:
        return "winograd";
    }
    return "";
}

model::PortMemoryLayout CalculateMemoryLayout(int numRows, int numColumns, int numChannels, int padding)
{
    // Calculate dimension parameters
    model::MemoryShape size{ numRows, numColumns, numChannels };
    model::MemoryShape offset{ padding, padding, 0 };
    model::MemoryShape stride{ numRows + 2 * padding, numColumns + 2 * padding, numChannels };

    return { size, stride, offset };
}
} // namespace

//
// Timing functions
//

template <typename ValueType>
auto TimeReferenceConvolution(const math::ChannelColumnRowTensor<ValueType>& signal, const math::ChannelColumnRowTensor<ValueType>& filters, int numFilters, int numIterations, dsp::ConvolutionMethodOption algorithm)
{
    const auto filterSize = static_cast<int>(filters.NumColumns());

    // Perform the convolution
    utilities::MillisecondTimer timer;
    if (algorithm == dsp::ConvolutionMethodOption::winograd)
    {
        const auto order = dsp::WinogradFilterOrder::tilesFirst;
        const int tileSize = 2;
        auto transformedFilters = dsp::GetTransformedFilters(filters, numFilters, tileSize, order);
        timer.Reset();
        for (int iter = 0; iter < numIterations; ++iter)
        {
            volatile auto result = Convolve2DWinogradPretransformed(signal, transformedFilters, numFilters, tileSize, filterSize, order);
        }
    }
    else
    {
        for (int iter = 0; iter < numIterations; ++iter)
        {
            volatile auto result = Convolve2D(signal, filters, static_cast<int>(numFilters), algorithm);
        }
    }
    auto duration = timer.Elapsed();
    return duration;
}

template <typename ValueType>
static void TimeConvolutionNode(ImageShape inputShape, FiltersShape filterShape, int numIterations, dsp::ConvolutionMethodOption convolutionMethod, ConvolutionOptions options = {})
{
    using Tensor = math::ChannelColumnRowTensor<ValueType>;

    int inputRows = inputShape.numRows;
    int inputColumns = inputShape.numColumns;
    int numChannels = inputShape.numChannels;

    int numFilters = filterShape.numFilters;
    int filterSize = filterShape.numRows;
    assert(filterShape.numColumns == filterSize);
    int numFilterChannels = filterShape.numChannels == 0 ? numChannels : filterShape.numChannels; // "0" means "number of input channels"
    auto isDepthwiseSeparable = (numFilterChannels == 1) && (numChannels > 1);
    UNUSED(isDepthwiseSeparable);

    const int outputRows = inputRows;
    const int outputColumns = inputColumns;
    const int inputPadding = (filterSize - 1) / 2;
    const int outputPadding = 0;
    const int stride = 1;

    auto inputSize = (inputRows + 2 * inputPadding) * (inputColumns + 2 * inputPadding) * numChannels;
    auto data = std::vector<ValueType>(inputSize);

    auto filterWeightsSize = numFilters * filterSize * filterSize * numChannels;
    auto filter = std::vector<ValueType>(filterWeightsSize);

    model::Model model;
    auto inputNode = model.AddNode<model::InputNode<ValueType>>(inputSize);

    auto inputMemoryLayout = CalculateMemoryLayout(inputRows, inputColumns, numChannels, inputPadding);
    auto outputMemoryLayout = CalculateMemoryLayout(outputRows, outputColumns, numFilters, outputPadding);
    auto filterWeights = Tensor(numFilters * filterSize, filterSize, numFilterChannels, filter);

    const model::OutputPort<ValueType>* output = nullptr;
    switch (convolutionMethod)
    {
    case dsp::ConvolutionMethodOption::automatic:
        std::cout << "Testing 'automatic' method --- using 'simple' instead" << std::endl;
    // fallthrough
    case dsp::ConvolutionMethodOption::simple:
        output = &nodes::SimpleConvolution(inputNode->output, inputMemoryLayout, outputMemoryLayout, filterWeights, stride);
        break;
    case dsp::ConvolutionMethodOption::diagonal:
        output = &nodes::DiagonalConvolution(inputNode->output, inputMemoryLayout, outputMemoryLayout, filterWeights, stride);
        break;
    case dsp::ConvolutionMethodOption::unrolled:
        output = &nodes::UnrolledConvolution(inputNode->output, inputMemoryLayout, outputMemoryLayout, filterWeights, stride);
        break;
    case dsp::ConvolutionMethodOption::winograd:
        output = &nodes::WinogradConvolution(inputNode->output, inputMemoryLayout, outputMemoryLayout, filterWeights, stride, options.winogradOptions.tileSize, options.winogradOptions.filterOrder);
        break;
    }

    auto map = model::Map(model, { { "input", inputNode } }, { { "output", *output } });

    auto rawDataTensor = Tensor(inputRows, inputColumns, numChannels, data);
    auto paddedDataTensor = Tensor(inputRows + 2, inputColumns + 2, numChannels);
    paddedDataTensor.Fill(0);

    auto dataTensorReference = paddedDataTensor.GetSubTensor(inputPadding, inputPadding, 0, inputRows, inputColumns, numChannels);
    dataTensorReference.CopyFrom(rawDataTensor);
    auto paddedDataArray = paddedDataTensor.ToArray();

    utilities::MillisecondTimer timer;

    model::MapCompilerOptions settings;
    settings.compilerSettings.optimize = true;
    settings.compilerSettings.useBlas = true;
    settings.compilerSettings.parallelize = false;
    settings.verifyJittedModule = true;
    model::ModelOptimizerOptions optimizerOptions;
    model::IRMapCompiler compiler(settings, optimizerOptions);

    timer.Reset();
    auto compiledMap = compiler.Compile(map);
    auto compilationTime = timer.Elapsed();
    std::cout << "Time to compile model: " << compilationTime << " ms\n";

    timer.Reset();
    for (int index = 0; index < numIterations; ++index)
    {
        compiledMap.SetInputValue(0, paddedDataArray);
        volatile auto compiledResult = compiledMap.ComputeOutput<ValueType>(0);
    }
    auto compiledTime = timer.Elapsed();

    // Compare against reference version from dsp library
    auto referenceTime = TimeReferenceConvolution(paddedDataTensor, filterWeights, numFilters, numIterations, convolutionMethod);

    auto algName = GetConvAlgName(convolutionMethod);
    std::cout << "Total time for " << numIterations << " iterations of " << inputRows << " x " << inputColumns << " x " << numChannels << " -> " << numFilters << " " << algName << " convolutions: " << compiledTime << " ms\t"
              << "(reference: " << referenceTime << " ms)\n";
}

//
// Main driver function to call all the timing functions
//
void TimeDSPNodes()
{
    //
    // Timings on jitted models
    //
    TimeConvolutionNode<float>({ 240, 240, 3 }, { 16, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::simple);
    TimeConvolutionNode<float>({ 240, 240, 3 }, { 16, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::unrolled);
    TimeConvolutionNode<float>({ 240, 240, 3 }, { 16, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    std::cout << std::endl;

    TimeConvolutionNode<float>({ 100, 100, 16 }, { 32, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::simple);
    TimeConvolutionNode<float>({ 100, 100, 16 }, { 32, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::unrolled);
    TimeConvolutionNode<float>({ 100, 100, 16 }, { 32, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    std::cout << std::endl;

    TimeConvolutionNode<float>({ 32, 48, 64 }, { 256, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::simple);
    TimeConvolutionNode<float>({ 32, 48, 64 }, { 256, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::unrolled);
    TimeConvolutionNode<float>({ 32, 48, 64 }, { 256, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    std::cout << std::endl;

    TimeConvolutionNode<float>({ 64, 64, 16 }, { 16, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::simple);
    TimeConvolutionNode<float>({ 64, 64, 16 }, { 16, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::unrolled);
    TimeConvolutionNode<float>({ 64, 64, 16 }, { 16, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    std::cout << std::endl;

    TimeConvolutionNode<float>({ 64, 64, 32 }, { 32, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::simple);
    TimeConvolutionNode<float>({ 64, 64, 32 }, { 32, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::unrolled);
    TimeConvolutionNode<float>({ 64, 64, 32 }, { 32, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    std::cout << std::endl;

    TimeConvolutionNode<float>({ 64, 64, 64 }, { 64, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::simple);
    TimeConvolutionNode<float>({ 64, 64, 64 }, { 64, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::unrolled);
    TimeConvolutionNode<float>({ 64, 64, 64 }, { 64, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    std::cout << std::endl;

    TimeConvolutionNode<float>({ 64, 64, 128 }, { 128, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::simple);
    TimeConvolutionNode<float>({ 64, 64, 128 }, { 128, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::unrolled);
    TimeConvolutionNode<float>({ 64, 64, 128 }, { 128, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    std::cout << std::endl;

    // Winograd-specific stuff
    TimeConvolutionNode<float>({ 127, 127, 8 }, { 8, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 127, 127, 16 }, { 16, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 127, 127, 32 }, { 32, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 127, 127, 64 }, { 64, 3, 3, 0 }, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    // TimeConvolutionNode<float>({127, 127, 128}, {128, 3, 3, 0}, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    // TimeConvolutionNode<float>({127, 127, 256}, {256, 3, 3, 0}, 10, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });

    // Winograd-specific
    std::cout << "Regular (3D)\n";
    std::cout << "Tiles-first\n";
    TimeConvolutionNode<float>({ 64, 64, 1 }, { 1, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 64, 64, 2 }, { 2, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 64, 64, 4 }, { 4, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 64, 64, 8 }, { 8, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 127, 127, 1 }, { 1, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 127, 127, 2 }, { 2, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 127, 127, 4 }, { 4, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });
    TimeConvolutionNode<float>({ 127, 127, 8 }, { 8, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::tilesFirst });

    std::cout << "\n";
    std::cout << "Filters-first\n";
    TimeConvolutionNode<float>({ 64, 64, 1 }, { 1, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 2 }, { 2, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 4 }, { 4, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 8 }, { 8, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 1 }, { 1, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 2 }, { 2, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 4 }, { 4, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 8 }, { 8, 3, 3, 0 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });

    std::cout << "\n";
    std::cout << "Depthwise-separable\n";
    std::cout << "Filters-first\n";
    TimeConvolutionNode<float>({ 64, 64, 1 }, { 1, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 2 }, { 2, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 4 }, { 4, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 8 }, { 8, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 16 }, { 16, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 64, 64, 32 }, { 32, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 1 }, { 1, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 2 }, { 2, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 4 }, { 4, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 8 }, { 8, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 16 }, { 16, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
    TimeConvolutionNode<float>({ 127, 127, 32 }, { 32, 3, 3, 1 }, 100, dsp::ConvolutionMethodOption::winograd, { 2, dsp::WinogradFilterOrder::filtersFirst });
}
